
pub mod stream;
pub mod server;
pub mod client;

use crate::client::Client;
use crate::server::Server;
use std::collections::HashSet;

fn main() {
    // 0-client 1-server
    let mut mode: u32 = 0;
    let mut host: Option<String> = None;
    let mut port: Option<u16> = None;
    let mut proxy_ports: HashSet<(u16, u16)> = HashSet::new();
    let mut server_listen_host: String = String::from("127.0.0.1");

    let mut args = std::env::args();
    args.next();
    loop {
        let arg = match args.next() {
            Some(a) => a,
            None => { break; }
        };
        if arg.starts_with('-') {
            if arg.len() < 2 {
                eprintln!("Unknown switch '-'");
                std::process::exit(-1);
            }
            match arg.chars().nth(1).unwrap() {
                'l' => {
                    mode = 1;
                }
                'p' => {
                    let arg1 = if arg.len() == 2 {
                        match args.next() {
                            Some(a) => a,
                            None => {
                                eprintln!("Switch '-p' must followed by a number");
                                std::process::exit(-1);
                            }
                        }
                    } else {
                        String::from(&arg[2..])
                    };
                    if arg1.contains(":") {
                        let div = arg1.split(":");
                        let mut arg = Vec::<String>::new();
                        for a in div {
                            arg.push(String::from(a));
                        }
                        if arg.len() < 2 {
                            eprintln!("Port format: <port> or <client>:<server>");
                            std::process::exit(-1);
                        }
                        let port1 = match arg.get(0).unwrap().parse::<u16>() {
                            Ok(p) => p,
                            Err(_) => {
                                eprintln!("Port cannot be '{}' but a number", arg1);
                                std::process::exit(-1);
                            }
                        };
                        let port2 = match arg.get(1).unwrap().parse::<u16>() {
                            Ok(p) => p,
                            Err(_) => {
                                eprintln!("Port cannot be '{}' but a number", arg1);
                                std::process::exit(-1);
                            }
                        };
                        proxy_ports.insert((port1, port2));
                    }
                    else {
                        let port = match arg1.parse::<u16>() {
                            Ok(p) => p,
                            Err(_) => {
                                eprintln!("Port cannot be '{}' but a number", arg1);
                                std::process::exit(-1);
                            }
                        };
                        proxy_ports.insert((port, port));
                    }
                }
                'h' => {
                    let arg1 = if arg.len() == 2 {
                        match args.next() {
                            Some(a) => a,
                            None => {
                                eprintln!("Switch '-h' must followed by a host address");
                                std::process::exit(-1);
                            }
                        }
                    } else {
                        String::from(&arg[2..])
                    };
                    server_listen_host = arg1;
                }
                _ => {
                    eprintln!("Unknown switch '{}'", arg);
                    std::process::exit(-1);
                }
            }
        } else {
            match host {
                Some(_) => {
                    let p = match arg.parse::<u16>() {
                        Ok(p) => p,
                        Err(_) => {
                            eprintln!("Port must be a number");
                            std::process::exit(-1);
                        }
                    };
                    port = Some(p);
                }
                None => {
                    host = Some(arg);
                }
            }
        }
    }

    let h = match host {
        Some(h) => h,
        None => {
            eprintln!("Hostname must be given");
            std::process::exit(-1);
        }
    };
    let p = match port {
        Some(p) => p,
        None => {
            eprintln!("Connect port must be given");
            std::process::exit(-1);
        }
    };

    if mode != 1 && server_listen_host != "127.0.0.1" {
        eprintln!("\'-h\' switch can only be used in server mode");
        std::process::exit(-1);
    }

    match mode {
        0 => {
            let mut client = Client::new(h.as_str(), p, &proxy_ports);
            client.run();
        }
        1 => {
            let mut server = Server::new(h.as_str(), p,
                                         server_listen_host.as_str());
            server.run();
        }
        _ => { std::process::exit(-1); }
    }
}
